from flask_httpauth import HTTPBasicAuth
from flask import Flask, request
import threading
import os
import socket
import demjson
import requests
import datetime
import psutil
from time import sleep


class MonitorClass():
    def __init__(self):
        if os.path.exists("db/") == False:
            os.mkdir("db")
        if os.path.exists("static/") == False:
            os.mkdir("static")
        self.unLock()
        self.console_api_url = open("server.txt").read()
        threading.Thread(target=self.set_server_info).start()
        threading.Thread(target=self.report).start()


    def isLock(self):
        return os.path.exists("db/LOCK")

    def Lock(self):
        self.save_log("db/LOCK", 'w', self.get_now_datetime())

    def unLock(self):
        try:
            os.remove("db/LOCK")
        except:
            pass

    def save_log(self, filename, mode, content):
        '''一句话保存文件'''
        f = open(filename, mode)
        f.write(content)
        f.close()

    def report(self):
        '''上报状态'''
        while True:
            try:
                if os.path.exists("static/task.json"):
                    task_status = demjson.decode(open("static/task.json").read())
                else:
                    task_status = ""

                report_info = {
                    "server_status": self.get_server_info(),
                    "task_status": task_status
                }

                response = requests.post(self.console_api_url + "/submit_report", data={
                    "report_info": demjson.encode(report_info)
                })

                if response.text.find("success") > -1:
                    self.save_log("static/log.txt", 'a+', "{}  Report Success......\n".format(self.get_now_datetime()))
                else:
                    self.save_log("static/log.txt", 'a+',
                                  "{} Report Failed: {}\n".format(self.get_now_datetime(), response.text))

            except Exception as e:
                self.save_log("static/log.txt", 'a+', "{}  Report Exception : {}\n".format(self.get_now_datetime(), e))

            sleep(3)

    def bytes2human(self, n):
        '字节转换为其他单位返回为str'
        symbols = ('K', 'M', 'G', 'T', 'P', 'E', 'Z', 'Y')
        prefix = {}
        for i, s in enumerate(symbols):
            prefix[s] = 1 << (i + 1) * 10
        for s in reversed(symbols):
            if n >= prefix[s]:
                value = float(n) / prefix[s]
                return '%.2f %s' % (value, s)
        return '%.2f B' % (n)

    def set_server_info(self):
        '''设置全局信息'''
        while True:
            mem = psutil.virtual_memory().percent
            cpu = psutil.cpu_percent(interval=1)
            tot_before = psutil.net_io_counters()
            os.system("netstat -n | awk '/^tcp/ {++S[$NF]} END {for(a in S) print a, S[a]}' > static/conn.txt")
            sleep(1)
            tot_after = psutil.net_io_counters()
            uploadSpeed = self.bytes2human(tot_after.bytes_sent - tot_before.bytes_sent)
            downloadSpeed = self.bytes2human(tot_after.bytes_recv - tot_before.bytes_recv)
            self.save_log('static/server_status.json', 'w', demjson.encode({
                'mem': mem,
                'cpu': cpu,
                'up': uploadSpeed,
                'down': downloadSpeed
            }))

    def get_server_info(self):
        '''获取全局信息'''
        try:
            server_status_file = open("static/server_status.json", 'r')
            server_status = demjson.decode(server_status_file.read())
            server_status_file.close()
            return server_status
        except:
            return {
                'mem': 0,
                'cpu': 0,
                'up': 0,
                'down': 0
            }

    def get_now_datetime(self):
        return datetime.datetime.strftime(datetime.datetime.now(), '%Y-%m-%d %H:%M:%S')


app = Flask(__name__)
Monitor = MonitorClass()
users = [
    {'username': '0x007er', 'password': 'admins'}
]

auth = HTTPBasicAuth()


@auth.get_password
def get_password(username):
    for user in users:
        if user['username'] == username:
            return user['password']
    return None


@app.route('/')
def hello_world():
    return "hello world"


@app.route('/submit', methods=['POST'])
def submit():
    '''
    提交任务
    taskinfo = {
        "type": "http",
        "charset": "utf-8",
        "allow_redirects": "0",
        "contain_list": ["hello", "world"],
        "exclude_list": ["cfuid", "acw"],
        "headers": {"host": "xss.app"},
        "path": "/login",
        "method": "GET",
        "thread": 3,
        "timeout": 3,
        "ipfile_url": "http://www.qq.com/1.txt",
        "task_id": "dev-test-1",
        "force":1
    }
    {
	"allow_redirects": "0",
	"charset": "utf-8",
	"contain_list": ["ezXSS", "POST"],
	"exclude_list": ["cfuid", "acw"],
	"force": 1,
	"headers": {
		"host": "xss.app"
	},
	"ipfile": "ip.txt",
	"method": "GET",
	"path": "/login",
	"task_id": "dev-test-1",
	"thread": 30,
	"timeout": 3,
	"type": "http"
}
    '''

    Monitor.save_log("taskinfo.json", 'w', request.form['task_info'])
    if Monitor.isLock() == True:
        return demjson.encode({
            "status": "failed",
            "text": "server is not ready"
        })
    else:
        os.system("nohup pypy3 Detector.py > static/runlog.txt 2>&1 &")
        return demjson.encode({
            "status": "success",
            "text": "submit success"
        })


@app.route('/stop')
def stop():
    '''
    停止任务
    :return:
    '''
    os.system("ps -ef|grep Detector |grep -v grep|cut -c 9-15|xargs kill -9")
    Monitor.unLock()
    return demjson.encode({
        "status": "success",
        "text": "stop success"
    })

if __name__ == '__main__':
    app.debug = False
    app.run(host="0.0.0.0", port=8088)
